package utils

import chisel3.RawModule
import chisel3.stage.ChiselGeneratorAnnotation
import utils.GlobalConfigLoader.{GlobalConfig, config}
import utils.module.NegativeResetWrapper

import scala.collection.mutable.{ArrayBuffer, ListBuffer}

object generate {
  def doBeforeGenerate(): Unit = {
    config.RUN_NAMING_WIRE = true
    BasicLogger.colorLog = false
    BasicLogger.filePathLog = false
  }

  def doAfterGenerate(): Unit = {
    config.RUN_NAMING_WIRE = false
    BasicLogger.colorLog = true
    BasicLogger.filePathLog = true
  }

  def doBeforeSynth(): Unit = {
    logger.info("Enter Synth mode")
    GlobalConfig.synthMode = true
  }

  def doAfterSynth(): Unit = {
    GlobalConfig.synthMode = false
  }

  def generateVerilog(
                       generativeClasses: Map[String, () => RawModule],
                       args: Array[String] = Array(),
                       outsideModules: Boolean = true,
                       generateWrapper: Boolean = true,
                       synthMode: Boolean = false
                     ): Unit = {
    doBeforeGenerate()
    // logger.info(f"Core: Start PC = 0${config.core.startPC}; L2Cake: ${config.l2Cache.sizeBytes} Bytes")
    val toGenerateList = args.toList.filter(generativeClasses.contains)
    val toGenerate = (if (toGenerateList.isEmpty) generativeClasses.keys else toGenerateList)
      .filter(x => if (outsideModules) true else config.has(s"XM_${x.toUpperCase}") || x.endsWith("Wrapper"))
    // did not set -td and generate all modules, dynamic change -td arg
    val targetDirParent = if (args.contains("-tdp")) args(args.indexOf("-tdp") + 1) else "build"

    logger.info(s"generating Verilog: $toGenerate")
    val generatedModules = ListBuffer[String]()
    toGenerate.foreach(moduleName => {
      if (moduleName.contains("Synth") || synthMode) doBeforeSynth()
      val newArgs = ArrayBuffer[String]()
      args.filter(a => a != "-tdp" && a != targetDirParent).foreach(newArgs.append)
      if (!newArgs.contains("-td")) newArgs ++= List("-td", f"$targetDirParent/chisel-tmp/$moduleName")

      def getNewOutputArgs(name: String = moduleName) = {
        val newArgsCopyInner = newArgs.clone()
        if (!newArgsCopyInner.contains("-o")) newArgsCopyInner ++= List("-o", f"../../chisel-rtl/$name.v")
        newArgsCopyInner
      }

      val newArgsCopy = getNewOutputArgs()
      logger.info(s"Module name: $moduleName; args: $newArgsCopy")
      (new chisel3.stage.ChiselStage).execute(newArgsCopy.toArray, Seq(ChiselGeneratorAnnotation(generativeClasses(moduleName))))
      var generateOkay = true
      try {
        if (generateWrapper) {
          val newArgsWrapperCopy = getNewOutputArgs(s"${moduleName}Wrapper")
          logger.info(s"Module name: ${moduleName}Wrapper; args: $newArgsWrapperCopy")
          (new chisel3.stage.ChiselStage).execute(newArgsWrapperCopy.toArray, Seq(ChiselGeneratorAnnotation(() =>
            new NegativeResetWrapper(generativeClasses(moduleName)(), moduleName = moduleName + "Wrapper"))))
        }
      } catch {
        case e: Exception =>
          logger.error(msg = s"Module ${moduleName}Wrapper generated failed! $e")
          e.printStackTrace()
          generateOkay = false
      }
      if (generateOkay) generatedModules += moduleName
      if (moduleName.contains("Synth") || synthMode) doAfterSynth()
    })
    logger.info(s"generated modules are: $generatedModules")
    logger.info(s"next run: make synth MODULES=\"${generatedModules.mkString(" ")}\"")
    doAfterGenerate()
  }

  def generateModule(name: String, module: () => RawModule) = generateVerilog(Map(name -> module))
}
